#include "common.hpp"

int main()
{
    // 设备资源初始化

    int32_t devices[] = {0, 1, 2, 3, 4, 5, 6, 7};           // 定义参与集合通信的卡的ID列表
    uint64_t data_count = (uint64_t)2 * GB / sizeof(float); // 参与集合通信的数据量, 按照fp32的个数计算
    uint64_t n_loop = 100;

    uint32_t ndev = sizeof(devices) / sizeof(devices[0]);
    HcclComm comms[ndev];
    ACLCHECK(aclInit(NULL));
    for (int32_t i = 0; i < ndev; i++)
        ACLCHECK(aclrtSetDevice(devices[i]));

    // 初始化通信域
    HCCLCHECK(HcclCommInitAll(ndev, devices, comms));

    auto inner_func = [data_count, n_loop](void *arg) -> int
    {
        ThreadContext *ctx = (ThreadContext *)arg;

        //  申请通信用device、sendBuf，recvBuf内存、stream等资源
        void *sendBuff;
        void *recvBuff;
        void *hostBuff;
        aclrtStream stream;
        uint64_t mallocSize = data_count * sizeof(float);

        ACLCHECK(aclrtSetDevice(ctx->device));
        ACLCHECK(aclrtCreateStream(&stream));

        // 初始化输入内存
        ACLCHECK(aclrtMalloc((void **)&sendBuff, mallocSize, ACL_MEM_MALLOC_HUGE_FIRST));
        ACLCHECK(aclrtMalloc((void **)&recvBuff, mallocSize, ACL_MEM_MALLOC_HUGE_FIRST));
        // 执行AllReduce操作
        for (int i = 0; i < n_loop; i++)
            HCCLCHECK(HcclAllReduce((void *)sendBuff, (void *)recvBuff, data_count, HCCL_DATA_TYPE_FP32, HCCL_REDUCE_SUM, ctx->comm, stream));
        ACLCHECK(aclrtSynchronizeStream(stream));

        // 释放通信用sendBuf、recvBuf内存，stream等资源
        ACLCHECK(aclrtFree(recvBuff));
        ACLCHECK(aclrtFree(sendBuff));
        ACLCHECK(aclrtDestroyStream(stream));
        ACLCHECK(aclrtResetDevice(ctx->device));
        return 0;
    };
    // 启动线程执行集合通信操作
    std::vector<std::unique_ptr<std::thread>>
        threads(ndev);
    struct ThreadContext args[ndev];
    for (uint32_t i = 0; i < ndev; i++)
    {
        args[i].device = i;
        args[i].ndev = ndev;
        args[i].comm = comms[i];
        threads[i].reset(new (std::nothrow) std::thread(inner_func, (void *)&args[i]));
    }
    for (int32_t i = 0; i < ndev; i++)
        threads[i]->join();
    for (uint32_t i = 0; i < ndev; i++)
        HCCLCHECK(HcclCommDestroy(comms[i]));
    ACLCHECK(aclFinalize());

    return 0;
}